#include "AM.h"
#include "printf.h"

#define NAIVE_MODE			0x89
#define XOR_MODE			0x90

#define CLIENT_1_ID			10
#define CLIENT_2_ID 		11
#define BUFFER_SIZE			12

#define MAX_PAYLOAD_SIZE	28
#define MAX_RETRIES			5
#define RETRY_DELAY			15

module MoteHubC @safe()
{
	uses {
		interface Boot;
		interface PacketLink;
		interface Leds;
		interface SplitControl as RadioControl;

		interface AMSend as RadioSend[am_id_t id];
		interface Receive as RadioReceive[am_id_t id];
		interface Packet as RadioPacket;
		interface AMPacket as RadioAMPacket;
	}
}
implementation
{
	// TODO: PHN - Move to external library
	typedef nx_struct FileMessage {
		nx_uint16_t packetId;
		nx_uint8_t data[MAX_PAYLOAD_SIZE-2];
	} FileMessage;

	void toggleRx();
	void toggleTx();
	void toggleEx();
	message_t* receive(message_t *msg, void *payload, uint8_t len);
	int addr2Client(am_addr_t addr);
	am_addr_t client2Addr(int client);
	task void handleMessages();
	task void unicastTask();
	task void broadcastTask();

	int transferMode = 0;
	am_addr_t localAddress;

	bool messageTaskBusy = FALSE;
	message_t buffer[2][BUFFER_SIZE];
	int bufferIndex[2];
	int bufferLength[2];
	
	bool unicastTaskBusy = FALSE;
	message_t unicastBuffer[BUFFER_SIZE];
	int unicastBufferIndex;
	int unicastBufferLength;
	
	bool broadcastTaskBusy = FALSE;
	message_t broadcastBuffer[BUFFER_SIZE];
	int broadcastBufferIndex;
	int broadcastBufferLength;

	event void Boot.booted(){
		bufferIndex[0] = bufferIndex[1] = 0;
		bufferLength[0] = bufferLength[1] = 0;
	
		printf("BOOTING\n"); printfflush();
		call RadioControl.start();		
	}
 
	event void RadioControl.stopDone(error_t error){ }

	event void RadioControl.startDone(error_t error){
		localAddress = (call RadioAMPacket.address());
	}

	event void RadioSend.sendDone[am_id_t id](message_t *msg, error_t error) {
		if(transferMode == NAIVE_MODE) {
			post unicastTask();
		} else {
			post broadcastTask();
		}
		post handleMessages();
	}

	event message_t * RadioReceive.receive[am_id_t id](message_t *msg, void *payload, uint8_t len){
		return receive(msg, payload, len);
	}
	
	message_t* receive(message_t *msg, void *payload, uint8_t len) {

		message_t *ret = msg;
		int client;
		int newBufferIndex;
		am_id_t type = (call RadioAMPacket.type(msg));

		// Only handle message sent directly to hub
		if(localAddress != (call RadioAMPacket.destination(msg)) || (type != NAIVE_MODE && type != XOR_MODE)) {
			return ret;
		}

		// Determine which client sent message
		client = addr2Client(call RadioAMPacket.source(msg));
		if(client < 0) {
			toggleEx();
			printf("UNKNOWN CLIENT\n"); printfflush();
			return ret;
		}
	
		// Detect overflow
		if(bufferLength[client] == BUFFER_SIZE) {
			toggleEx(); // Overflow
			printf("BUFFER %d OVERFLOW\n", client); printfflush();
			return msg;
		}

		// Place message in buffer
		newBufferIndex = (bufferIndex[client] + bufferLength[client]) % BUFFER_SIZE;
		buffer[client][newBufferIndex] = *msg;
		bufferLength[client]++;

		if(!messageTaskBusy) {
			post handleMessages();
			messageTaskBusy = TRUE;
		}
	
		toggleRx();				
		return ret;
	}

	task void handleMessages() {
	
		message_t* msg[2];
		uint8_t len[2];
		FileMessage* payload[2];
		int i;

		message_t xorMsg;
		FileMessage* newPayload;

		// Only send if both buffers contain messages
		if(!(bufferLength[0] > 0) || !(bufferLength[1] > 0)) {
			//printf("MAIN BUFFERS EMPTY\n"); printfflush();
			messageTaskBusy = FALSE;
			return;
		}
	
		msg[0] = &buffer[0][bufferIndex[0]];
		msg[1] = &buffer[1][bufferIndex[1]];
		len[0] = call RadioPacket.payloadLength(msg[0]); 
		len[1] = call RadioPacket.payloadLength(msg[1]); 
		payload[0] = call RadioPacket.getPayload(msg[0], len[0]);
		payload[1] = call RadioPacket.getPayload(msg[1], len[1]);

		// Start new session on initial message
		if(payload[0]->packetId == 1 || payload[1]->packetId == 1) {
			if(payload[0]->packetId == 1 && payload[1]->packetId == 1) {
				transferMode = (call RadioAMPacket.type(msg[0]));
				printf("NEW SESSION\n"); printfflush();
			} else {
				toggleEx();
				printf("SESSION INIT ERROR\n"); printfflush();
				return;
			}
		}

		if(transferMode == NAIVE_MODE) {

			// Detect overflow
			if(unicastBufferLength+1 >= BUFFER_SIZE) {
				toggleEx(); // Overflow
				printf("UNICAST BUFFER OVERFLOW\n"); printfflush();
				return;
			}
	
			// Place message in buffer
			i = (unicastBufferIndex + unicastBufferLength) % BUFFER_SIZE;
			unicastBuffer[i] = *msg[0];
			unicastBuffer[i+1] = *msg[1];
			unicastBufferLength += 2;
	
			bufferIndex[0] = (bufferIndex[0] + 1) % BUFFER_SIZE;
			bufferIndex[1] = (bufferIndex[0] + 1) % BUFFER_SIZE;
			bufferLength[0]--;
			bufferLength[1]--;
	
			if(!unicastTaskBusy) {
				post unicastTask();
				unicastTaskBusy = TRUE;
			} else {
				post handleMessages();
			}
	
	
		} else { // XOR_MODE

			// Detect overflow
			if(broadcastBufferLength+1 >= BUFFER_SIZE) {
				toggleEx(); // Overflow
				printf("BROADCAST BUFFER OVERFLOW\n"); printfflush();
				return;
			}
	
			xorMsg = *msg[0];
			newPayload = (FileMessage *) (call RadioPacket.getPayload(&xorMsg, len[0]));
	
			// TODO: PHN - Should first byte be ignored (contains packet ID, should be same, and will then always XOR to 0?)
			printf("SEQ# BEFORE: %d\n", newPayload->packetId); printfflush();
			for(i = 0; i < (len[0]-2); ++i) {
				newPayload->data[i] = payload[0]->data[i] ^ payload[1]->data[i];
			} 
			printf("SEQ# AFTER: %d\n", newPayload->packetId); printfflush();
			printf("%d %d", newPayload->data[0], newPayload->data[1]); printfflush();
	
			// Place message in buffer
			i = (broadcastBufferIndex + broadcastBufferLength) % BUFFER_SIZE;
			broadcastBuffer[i] = xorMsg;
			broadcastBufferLength++;
	
			bufferIndex[0] = (bufferIndex[0] + 1) % BUFFER_SIZE;
			bufferIndex[1] = (bufferIndex[0] + 1) % BUFFER_SIZE;
			bufferLength[0]--;
			bufferLength[1]--;
	
			if(!broadcastTaskBusy) {
				post broadcastTask();
				broadcastTaskBusy = TRUE;
			} else {
				post handleMessages();
			}
		}	
	}

	void task unicastTask() {

		message_t* msg;
		uint8_t len;
		am_addr_t dest;

		if(unicastBufferLength < 1) {
			unicastTaskBusy = FALSE;
			return;
		}

		msg = &unicastBuffer[unicastBufferIndex];
		len = call RadioPacket.payloadLength(msg); 
	
		call PacketLink.setRetries(msg, MAX_RETRIES);
		call PacketLink.setRetryDelay(msg, RETRY_DELAY);

		dest = ((call RadioAMPacket.source(msg)) == CLIENT_1_ID)? CLIENT_2_ID : CLIENT_1_ID;
	
		if((call RadioSend.send[NAIVE_MODE](dest, msg, len)) != SUCCESS) {						
			toggleEx();
			printf("FAILED TO SEND\n"); printfflush();
			post unicastTask();			
		} else {
			toggleTx();
			unicastBufferIndex = (unicastBufferIndex + 1) % BUFFER_SIZE;
			unicastBufferLength--;			
			printf("SENT FROM %d to %d\n", call RadioAMPacket.source(msg), call RadioAMPacket.destination(msg)); printfflush();			
		}
	}
	
	void task broadcastTask() {

		message_t* msg;
		uint8_t len;

		if(broadcastBufferLength < 1) {
			broadcastTaskBusy = FALSE;
			return;
		}

		msg = &broadcastBuffer[broadcastBufferIndex];
		len = call RadioPacket.payloadLength(msg); 
	
		call PacketLink.setRetries(msg, MAX_RETRIES);
		call PacketLink.setRetryDelay(msg, RETRY_DELAY);

		if((call RadioSend.send[XOR_MODE](AM_BROADCAST_ADDR, msg, len)) != SUCCESS) {						
			toggleEx();
			printf("FAILED TO BROADCAST\n"); printfflush();
			post broadcastTask();			
		} else {
			toggleTx();
			broadcastBufferIndex = (broadcastBufferIndex + 1) % BUFFER_SIZE;
			broadcastBufferLength--;			
			printf("BROADCASTING FROM %d to %d\n", call RadioAMPacket.source(msg), call RadioAMPacket.destination(msg)); printfflush();			
		}
	} 
	
	void toggleRx() {
		call Leds.led1Toggle();
	}	

	void toggleTx() {
		call Leds.led0Toggle();
	}	

	void toggleEx() {
		call Leds.led2Toggle();
	}
	
	int addr2Client(am_addr_t addr) {
		switch (addr) {
			case CLIENT_1_ID:
			return 0;
			case CLIENT_2_ID:
			return 1;
			default:
			return -1;
		}		
	}

	am_addr_t client2Addr(int client) {
		switch (client) {
			case 0:
			return CLIENT_1_ID;
			case 1:
			return CLIENT_2_ID;
			default:
			return -1;
		}		
	}
}